import PIL
import numpy as np
from PIL import Image
from pathlib import Path


def random_walk(mu, x, sigma, N):
	i = 0
	while i < N:
		yield x
		w = np.random.normal(0, sigma)
		x = round(mu * x + w, 2)
		i += 1


def random_zip_main():
	N = 5
	walk1 = list(random_walk(0, 1, 1.5, N))
	walk2 = list(random_walk(0, 2, 0.5, N))
	walk3 = list(random_walk(0, 3, 2.5, N))

	for item in zip(walk1, walk2, walk3):
		print(item)


class FaceDataset:
	def __init__(self, path, start=0, stop=999999, step=1):
		self._path = Path(path)
		self._lis = list()
		self._img = None
		self._iterpos = start
		self._stop = stop
		self._step = step

	def load_pic(self):
		image = self._path.rglob('*.jpg')
		images = list(image)
		self._lis = images

	def trans_pic(self, image):
		img = np.array(Image.open(image))
		self._img = img
		return img

	def __iter__(self):
		return self

	def __next__(self):
		if self._iterpos <= self._stop:
			self.trans_pic(self._lis[self._iterpos])
			self._iterpos += self._step
			return self._img
		else:
			raise StopIteration("达到加载上限。")


def face_dataset_main():
	path = r'D:\buaa_czw\2022 Autumn\General Programming 2022\Week 9\originalPics'
	d = FaceDataset(path, stop=50)
	d.load_pic()
	for item in d:
		print(item)


if __name__ == "__main__":
	# random_zip_main()
	face_dataset_main()
